import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from dash import dcc, html
from dash.dependencies import Input, Output, State, ALL
import dash
from sympy import latex


class VectorSpaces:
    """
    向量空间可视化组件
    包含子空间、线性无关、基、维数等核心概念的交互式演示
    """
    
    def __init__(self):
        self.vectors = []
        self.basis = []
        self.pivot_indices = []
        self.matrix_dimensions = {'rows': 0, 'cols': 3}
        
    def get_component(self):
        """返回Dash组件"""
        return html.Div([
            html.H3('向量空间概念演示'),
            
            # 矩阵和向量输入控件
            html.Div([
                # 动态矩阵输入
                html.Div([
                    html.H4('动态矩阵A'),
                    html.Div(id='dynamic-matrix-inputs'),
                    dcc.Store(id='matrix-dimensions', data={'rows': 0, 'cols': 3}),
                    dcc.Store(id='vectors-store', data=[])
                ], style={'marginBottom': '20px'}),
                
                # 向量输入
                html.Div([
                    html.Label('新向量:', style={'marginRight': '10px'}),
                    dcc.Input(id='new-x', type='number', value=None, step=0.1, placeholder='x', style={'width': '70px', 'marginRight': '5px'}),
                    dcc.Input(id='new-y', type='number', value=None, step=0.1, placeholder='y', style={'width': '70px', 'marginRight': '5px'}),
                    dcc.Input(id='new-z', type='number', value=None, step=0.1, placeholder='z', style={'width': '70px', 'marginRight': '10px'}),
                    html.Button('➕ 添加', id='add-vector', style={'marginRight': '5px'}),
                    dcc.Dropdown(id='vector-to-remove', options=[], placeholder='选择要删除的向量', style={'width': '200px', 'marginRight': '5px'}),
                    html.Button('➖ 删除', id='remove-vector', style={'marginRight': '5px'}),
                    html.Button('✏️ 修改', id='modify-vector', style={'marginRight': '5px'}),
                    html.Button('🗑️ 清除', id='clear-vectors'),
                ], style={'display': 'flex', 'alignItems': 'center', 'margin': '20px'}),
                
                html.Div([
                    html.H4('数学推导过程', style={'borderBottom': '2px solid #666', 'paddingBottom': '5px'}),
                    dcc.Markdown(id='math-derivation', mathjax=True,
                        style={'whiteSpace': 'pre-wrap', 'padding': '15px',
                               'backgroundColor': '#f8f9fa', 'borderRadius': '8px'})
                ], style={'marginTop': '20px'}),
            ]),
            
            # 3D图形展示
            dcc.Graph(id='vector-space-plot'),
            
            # 结果显示
            html.Div(id='basis-matrix-result', style={'margin': '20px', 'padding': '10px', 'border': '1px solid #ddd'}),
            
            # 概念说明
            html.Div([
                html.H4('向量空间数学原理与应用'),
                
                html.Div([
                    html.H5('数学家贡献'),
                    html.P('• Hermann Grassmann (1809-1877): 首次提出向量空间概念'),
                    html.P('• Giuseppe Peano (1858-1932): 公理化定义向量空间'),
                ], style={'marginBottom': '10px'}),
                
                html.Div([
                    html.H5('数学原理'),
                    html.P('• 公理化定义: 满足加法和标量乘法的8条公理'),
                    html.P('• 子空间: 非空子集且对加法和标量乘法封闭'),
                    html.P('• 基与维数: 极大线性无关组称为基，基中向量个数为维数'),
                ], style={'marginBottom': '10px'}),
                
                html.Div([
                    html.H5('应用实例'),
                    html.P('• 计算机图形学: 3D模型的空间变换'),
                    html.P('• 机器学习: 特征空间、核方法'),
                    html.P('• 物理模拟: 力场和速度场的表示'),
                ]),
                
                html.Div(id='dynamic-explanation')
            ], style={'marginBottom': '20px', 'padding': '10px', 'border': '1px solid #ddd', 'borderRadius': '5px'})
        ])
        
    def register_callbacks(self, app):
        """注册回调函数"""
        
        # 动态生成矩阵输入行
        @app.callback(
            Output('dynamic-matrix-inputs', 'children'),
            Output('matrix-dimensions', 'data'),
            [
                Input('add-vector', 'n_clicks'),
                Input('remove-vector', 'n_clicks'),
                Input('clear-vectors', 'n_clicks'),
                Input('modify-vector', 'n_clicks')
            ],
            [
                State('matrix-dimensions', 'data'),
                State({'type': 'matrix-input', 'row': ALL, 'col': ALL}, 'value')
            ]
        )
        def update_matrix_inputs(add_clicks, remove_clicks, clear_clicks, modify_clicks, dimensions, matrix_values):
            rows = len(self.vectors)
            cols = 3  # 固定3列(x,y,z)
            
            # 生成动态输入行
            children = []
            for i in range(rows):
                row_inputs = []
                for j in range(cols):
                    row_inputs.append(
                        dcc.Input(
                            id={'type': 'matrix-input', 'row': i, 'col': j},
                            type='number',
                            value=self.vectors[i][j] if i < len(self.vectors) else 0,
                            style={'width': '70px', 'marginRight': '5px'}
                        )
                    )
                children.append(html.Div(row_inputs, style={'marginBottom': '5px'}))
            
            # 更新存储的维度信息
            current_dimensions = dimensions or {'rows': 0, 'cols': 3}
            
            if rows != current_dimensions['rows'] or cols != current_dimensions['cols']:
                return children, {'rows': rows, 'cols': cols}
            
            return dash.no_update, dash.no_update
            # 原子化更新矩阵和向量存储
            with self.app.server.app_context():
                if matrix_values and len(matrix_values) == rows*cols:
                    self.vectors = np.array(matrix_values).reshape(rows, cols).tolist()
                    dash.callback_context.stored_values['vectors-store'] = self.vectors
                else:
                    matrix_values = [v for vec in self.vectors for v in vec]
                    dash.callback_context.stored_values['matrix-input'] = matrix_values
            
            return children, current_dimensions
        
        @app.callback(
            Output('vector-to-remove', 'options'),
            Output('vectors-store', 'data'),
            Input('add-vector', 'n_clicks'),
            Input('remove-vector', 'n_clicks'),
            Input('clear-vectors', 'n_clicks'),
            Input('modify-vector', 'n_clicks'),
            State('matrix-dimensions', 'data'),
            State('vector-to-remove', 'value'),
            State('new-x', 'value'),
            State('new-y', 'value'),
            State('new-z', 'value'),
            State('vectors-store', 'data'),
            prevent_initial_call=True)
        def update_vector_dropdown(add_clicks, remove_clicks, clear_clicks, modify_clicks, matrix_dimensions, selected_vector, new_x, new_y, new_z, store_data):
            ctx = dash.callback_context
            triggered_prop = ctx.triggered[0]['prop_id'].split('.')[0] if ctx.triggered else None
            # 新增输出参数
            matrix_dim_update = dash.no_update
            # 确保store_data正确初始化
            store_data = store_data or []
            if not hasattr(self, 'vectors') or self.vectors is None:
                self.vectors = store_data.copy()
            # 生成下拉框选项
            options = [{'label': f'向量{i+1}: ({vec[0]:.1f}, {vec[1]:.1f}, {vec[2]:.1f})', 'value': i} 
                      for i, vec in enumerate(self.vectors)]
            
            # 处理向量操作
            if ctx.triggered:
                trigger_id = ctx.triggered[0]['prop_id'].split('.')[0]
                
                if trigger_id == 'add-vector' and new_x is not None and new_y is not None and new_z is not None:
                    # 立即清除输入框值
                    new_x = None
                    new_y = None
                    new_z = None
                    # 强制更新matrix-dimensions
                    matrix_dim_update = {'rows': len(self.vectors)+1, 'cols': 3}
                    self.matrix_dimensions = matrix_dim_update
                    # 新增零向量检查
                    new_vec = [float(new_x or 0), float(new_y or 0), float(new_z or 0)]
                    if np.linalg.norm(new_vec) < 1e-8:
                        if triggered_prop == 'add-vector':
                            return options, self.vectors
                    # 确保只添加唯一向量
                    if not any(np.array_equal(new_vec, vec) for vec in self.vectors):
                        self.vectors.append(new_vec)
                        # 立即同步矩阵维度
                        self.matrix_dimensions['rows'] = len(self.vectors)
                    # 强制更新输入框状态
                    if triggered_prop == 'add-vector':
                        # 同步输入框状态
                        return options, self.vectors
                elif trigger_id == 'remove-vector' and selected_vector is not None:
                    if 0 <= selected_vector < len(self.vectors):
                        self.vectors.pop(selected_vector)
                        matrix_dim_update = {'rows': len(self.vectors), 'cols': 3}
                elif trigger_id == 'clear-vectors':
                    self.vectors = []
                    matrix_dim_update = {'rows': 0, 'cols': 3}
                    return options, store_data, matrix_dim_update
                elif trigger_id == 'add-vector':
                    # 同步输入框状态
                    return options, self.vectors, matrix_dim_update
                elif trigger_id == 'remove-vector' and selected_vector is not None:
                    if 0 <= selected_vector < len(self.vectors):
                        self.vectors.pop(selected_vector)
                        matrix_dim_update = {'rows': len(self.vectors), 'cols': 3}
                elif trigger_id == 'clear-vectors':
                    self.vectors = []
                    matrix_dim_update = {'rows': 0, 'cols': 3}
                    return options, store_data, matrix_dim_update
                elif trigger_id == 'clear-vectors':
                    self.vectors = []
                    self.basis = []
                    self.pivot_indices = []
                    matrix_A = np.empty((0, 3))  # 清空时重置矩阵维度
                
                # 同步到Store组件
                store_data = self.vectors.copy()
                
                # 调试输出
                print(f"更新后向量列表: {self.vectors}")
                print(f"存储数据: {store_data}")
                
                if triggered_prop == 'add-vector':
                    return options, store_data, matrix_dim_update
            return options, store_data

        @app.callback(
                Output('vector-space-plot', 'figure'),
                Output('basis-matrix-result', 'children'),
                Output('dynamic-explanation', 'children'),
                Output('math-derivation', 'children'),
                [Input('add-vector', 'n_clicks'),
                Input('remove-vector', 'n_clicks'),
                Input('clear-vectors', 'n_clicks'),
                Input('modify-vector', 'n_clicks'),
                Input({'type': 'matrix-input', 'row': ALL, 'col': ALL}, 'value'),
                Input('matrix-dimensions', 'data')],
                [State('new-x', 'value'),
                State('new-y', 'value'),
                State('new-z', 'value'),
                State('vector-to-remove', 'value')])
    
        def update_vector_space(add_clicks, remove_clicks, clear_clicks, modify_clicks,
                              matrix_inputs, matrix_dimensions,
                              new_x, new_y, new_z, vector_to_remove):
            ctx = dash.callback_context
            if not ctx.triggered:
                return go.Figure(), "", "", ""
                
            # 构建动态矩阵
            # 构建动态矩阵
            rows = len(self.vectors)
            cols = 3
            matrix_values = np.array(matrix_inputs).reshape(rows, cols) if matrix_inputs and rows*cols == len(matrix_inputs) else np.array([])
            matrix_A = matrix_values.astype(float) if matrix_values.size > 0 else np.empty((0, 3))
                
            # 初始化matrix_data字典
            matrix_data = {
                '原始向量空间': '',
                '基向量': '',
                '线性组合空间': '',
                '预留': ''
            }

            button_id = ctx.triggered[0]['prop_id'].split('.')[0]

            # 处理向量操作
            if button_id == 'add-vector' and new_x is not None and new_y is not None and new_z is not None:
                new_vector = np.array([new_x, new_y, new_z])
                self.vectors.append(new_vector)
            elif button_id == 'remove-vector' and len(self.vectors) > 0 and vector_to_remove is not None and 0 <= vector_to_remove < len(self.vectors):
                self.vectors.pop(vector_to_remove)
            elif button_id == 'modify-vector' and len(self.vectors) > 0 and vector_to_remove is not None and 0 <= vector_to_remove < len(self.vectors) and new_x is not None and new_y is not None and new_z is not None:
                self.vectors[vector_to_remove] = np.array([new_x, new_y, new_z])
            elif button_id == 'clear-vectors':
                self.vectors = []
                self.basis = []
                self.pivot_indices = []
                matrix_A = np.empty((0, 3))  # 清空时重置矩阵维度
            
            # 自动检查线性无关性（新增或删除向量时或矩阵变化时触发）
            if button_id in ['add-vector', 'remove-vector', 'clear-vectors'] or button_id.startswith('a') and len(self.vectors) >= 1:
                # 合并动态矩阵和向量时统一维度
                self.vectors_array = np.vstack(self.vectors) if len(self.vectors) > 0 else np.empty((0, 3))
                combined_matrix = np.vstack([matrix_A, self.vectors_array])
                
                rank = np.linalg.matrix_rank(combined_matrix)
                _, pivots = np.linalg.qr(combined_matrix.T)
                self.pivot_indices = np.where(pivots[:rank] != 0)[0]
                self.basis = [combined_matrix[i] for i in sorted(self.pivot_indices)]
                

            # 创建四宫格布局
            fig = make_subplots(
                rows=2, cols=2,
                specs=[[{'type': 'scene'}, {'type': 'scene'}],
                       [{'type': 'scene'}, {'type': 'scene'}]],
                subplot_titles=('矩阵', '基', '列空间', '预留'),
                vertical_spacing=0.15, horizontal_spacing=0.15
            )
            
            # 图1: 原始向量空间 - 独立向量彩色，相关向量黑色
            colors = ['blue', 'red', 'green', 'purple', 'orange', 'yellow']
            pivot_indices = getattr(self, 'pivot_indices', [])
            for i, vec in enumerate(self.vectors):
                color = colors[i % len(colors)] if i in pivot_indices else 'black'
                fig.add_trace(go.Scatter3d(
                    x=[0, vec[0]], y=[0, vec[1]], z=[0, vec[2]],
                    mode='lines+markers+text',
                    name=f'向量{i+1}',
                    text=[f'({vec[0]:.1f}, {vec[1]:.1f}, {vec[2]:.1f})'],
                    textposition='top center',
                    line=dict(color=color, width=3),
                    marker=dict(size=5)
                ), row=1, col=1)
            
            # 图2: 基向量
            if hasattr(self, 'basis') and self.basis:
                for i, vec in enumerate(self.basis):
                    fig.add_trace(go.Scatter3d(
                        x=[0, vec[0]], y=[0, vec[1]], z=[0, vec[2]],
                        mode='lines+markers',
                        name=f'基向量{i+1}',
                        line=dict(color=colors[i % len(colors)], width=5),
                        marker=dict(size=8)
                    ), row=1, col=2)
            
            # 图3: 列空间可视化
            if len(self.vectors) > 0:
                matrix = np.array(self.vectors)
                
                # 计算QR分解获取列空间
                Q, R = np.linalg.qr(matrix.T)
                rank = np.linalg.matrix_rank(matrix)
                
                # 显示列空间
                if rank == 1:
                    # 一维列空间(直线)
                    vec = Q[:,0]
                    t = np.linspace(-2, 2, 50)
                    x = t * vec[0]
                    y = t * vec[1]
                    z = t * vec[2]
                    
                    fig.add_trace(go.Scatter3d(
                        x=x, y=y, z=z,
                        mode='lines',
                        line=dict(color='purple', width=4),
                        name='列空间(直线)'
                    ), row=2, col=1)
                elif rank >= 2:
                    # 二维列空间(平面)
                    vec1 = Q[:,0]
                    vec2 = Q[:,1]
                    u = np.linspace(-1, 1, 10)
                    v = np.linspace(-1, 1, 10)
                    u, v = np.meshgrid(u, v)
                    
                    x = u * vec1[0] + v * vec2[0]
                    y = u * vec1[1] + v * vec2[1]
                    z = u * vec1[2] + v * vec2[2]
                    
                    fig.add_trace(go.Surface(
                        x=x, y=y, z=z,
                        opacity=0.5,
                        colorscale='Viridis',
                        showscale=False,
                        name='列空间(平面)'
                    ), row=2, col=1)
                    
                    # 更新矩阵数据
                    matrix_data['线性组合空间'] = f"列空间矩阵(Q):\n{Q[:,:rank]}"
            
            # 简化的布局配置
            axis_config = dict(range=[-3, 3])
            fig.update_layout(
                scene=dict(xaxis=axis_config, yaxis=axis_config, zaxis=axis_config),
                scene2=dict(xaxis=axis_config, yaxis=axis_config, zaxis=axis_config),
                scene3=dict(xaxis=axis_config, yaxis=axis_config, zaxis=axis_config),
                scene4=dict(xaxis=axis_config, yaxis=axis_config, zaxis=axis_config),                
                title_text='向量空间四宫格演示',
                showlegend=True,
                height=800,
                annotations=[
                    dict(
                        text=matrix_data['原始向量空间'],
                        x=0.2, y=0.9,
                        xref='paper', yref='paper',
                        showarrow=False,
                        font=dict(size=10)
                    ),
                    dict(
                        text=matrix_data['基向量'],
                        x=0.8, y=0.9,
                        xref='paper', yref='paper',
                        showarrow=False,
                        font=dict(size=10)
                    ),
                    dict(
                        text=matrix_data['线性组合空间'],
                        x=0.2, y=0.4,
                        xref='paper', yref='paper',
                        showarrow=False,
                        font=dict(size=10)
                    )
                ]
            )
            
            # 添加双击事件处理
            # 添加双击事件处理
            fig.update_layout(
                scene_dragmode='orbit',
                # 双击放大功能
                clickmode="event+select",
                scene=dict(
                    aspectmode='auto',
                    aspectratio=dict(x=1, y=1, z=1)
                ),
                # 双击事件处理
                uirevision='zoom'
            )
            
            # 初始化返回变量
            explanation = ''
            basis_matrix = ''
            derivation_content = ''
            original_matrix_content = ''
            
            # 为每个子图创建矩阵数据
            matrix_data = {
                '原始向量空间': '',
                '基向量': '',
                '线性组合空间': '',
                '预留': ''
            }
            
            if len(self.vectors) > 0:
                matrix = np.array(self.vectors)
                matrix_data['原始向量空间'] = f"原始矩阵:\n{matrix}"
                
                if hasattr(self, 'basis') and self.basis:
                    basis_matrix = np.array(self.basis)
                    matrix_data['基向量'] = f"基矩阵:\n{basis_matrix}"
                    
                    if len(self.basis) >= 2:
                        vec1 = self.basis[0]
                        vec2 = self.basis[1]
                        span_matrix = np.array([vec1, vec2])
                        matrix_data['线性组合空间'] = f"生成空间矩阵:\n{span_matrix}"
            
            if hasattr(self, 'basis') and self.basis:
                basis_matrix = html.Pre(
                    '基矩阵:\n' + 
                    '\n'.join([f'向量{i+1}: {vec}' for i, vec in enumerate(self.basis)])
                )
                explanation = html.Div([
                    html.H5('当前空间分析'),
                    html.P(f'空间维度: {len(self.basis)}'),
                    html.P('基向量如上所示')
                ])
            
            if len(self.vectors) > 0:
                matrix = np.array(self.vectors)
                rank = np.linalg.matrix_rank(matrix)
                
                # 生成QR分解推导
                q_matrix, r_matrix = np.linalg.qr(matrix.T)
                qr_derivation = (
                    f"$$\nQ = {latex(np.round(q_matrix, 2))}\nR = {latex(np.round(r_matrix, 2))}$$\n"
                    "**正交基Q矩阵**的列向量构成规范正交基,**上三角矩阵R**的对角线元素显示线性独立性"
                )
                
                # 生成秩计算推导
                rank_derivation = f"$$\n秩(A) = {rank}\n$$\n矩阵的秩表示**线性无关**的列向量个数，此处通过主元位置确定"
                
                # 生成基向量推导
                basis_derivation = "基向量选择过程：\n"
                basis_derivation += "\\n".join([f"1. 向量{i+1}: {vec} {'' if i in self.pivot_indices else '(线性相关)'}" 
                                                for i, vec in enumerate(self.vectors)])
                
                derivation_content = f"### QR分解过程\n{qr_derivation}\n\n### 秩计算\n{rank_derivation}\n\n### 基向量选择\n{basis_derivation}"
            
            if len(self.vectors) > 0:
                matrix = np.array(self.vectors)
                original_matrix_content = html.Pre(
                    '原始矩阵:\n' + 
                    '\n'.join([f'向量{i+1}: {vec}' for i, vec in enumerate(self.vectors)])
                )
            
            return fig, basis_matrix, explanation, derivation_content